#pragma once
#include <iostream>

namespace me
{
    enum class color
    {
        RED,
        BLACK
    };

    template <typename Value>
    struct rbtreeNode
    {
        Value _data;
        rbtreeNode* _left;
        rbtreeNode* _right;
        rbtreeNode* _parent;
        color _col;

        rbtreeNode(const Value& data)
            :_data(data), _left(nullptr), _right(nullptr), _parent(nullptr)
        {}
    };

    template <typename V, typename Ref, typename Ptr>
    struct rbtreeIterator
    {
        using node = rbtreeNode<V>;
        using Self = rbtreeIterator<V, Ref, Ptr>;
        node* _node;
        node* _root; // 用来处理end迭代器为空，以便找到最右节点

        rbtreeIterator(node* other, node* root):_node(other), _root(root) {}

        // 走到中序的下一个位置
        Self& operator++()
        {
            if(_node->_right) {
                // 右子树不为空，走到右子树的最左节点
                node* mostLeft = _node->_right;
                while(mostLeft->_left) {
                    mostLeft = mostLeft->_left;
                } 
                _node = mostLeft;
            }else {
                // 右子树为空，往上走，直到cur是parent的左子树
                node* cur = _node;
                node* parent = cur->_parent;
                // 当已经遍历到最后一个时，此时++向上返回cur一直是parent的右子树，当cur为根的时候，parent为空，此时结束
                while(parent && cur == parent->_right) {
                    cur = parent;
                    parent = cur->_parent;
                }
                _node = parent;
            }
            return *this;
        }

        Self& operator--()
        {
            // 中序遍历——左——根——右，operator--结合end实现反向遍历，即右——根——左

            if(_node == nullptr) {
                // 当_node == nullptr，意味着此时迭代器从end开始，我们要返回红黑树的最右节点
                node* mostRight = _root;
                while(mostRight && mostRight->_right) {
                    mostRight = mostRight->_right;
                }
                _node = mostRight;
            }else if(_node->_left) {
                // 左子树不为空，走到左子树的最右节点
                node* mostRight = _node->_left;
                while(mostRight->_right) {
                    mostRight = mostRight->_right;
                }
                _node = mostRight;
            }else {
                // 左子树为空，向上走，直到cur是parent的右子树
                node* cur = _node;
                node* parent = cur->_parent;
                while(parent && cur == parent->_left) {
                    cur = parent;
                    parent = cur->_parent;
                }
                _node = parent;
            } 
            return *this;
        }

        Ref operator*() { return _node->_data; }
        Ptr operator->() { return &(_node->_data); }
        bool operator==(const Self& other) const { return _node == other._node; }
        bool operator!=(const Self& other) const { return _node != other._node; }
    };

    // 第二个模板参数用来同一map和set，因为set只有一个key，而map有key/value
    // 第二个模板参数set就传key，而map传std::pair<key,value>
    // 这个模板参数才是红黑树节点底层真正存储的数据
    // 第三个模板参数用来取出value中的key值，因为不确定到底存储的是key/<key,value>，所以需要有上层控制，
    // 使用不同的容器，传入不同的keyofvalue，以此同一比较时的逻辑
    template <typename K, typename Value, typename KeyOfValue>
    class rbtree
    {
        using iterator = rbtreeIterator<Value, Value&, Value*>;
        using const_iterator = rbtreeIterator<Value, const Value&, const Value*>;
    public:
        // 迭代器相关接口

        /*返回最左节点*/
        iterator begin()
        {
            node* cur = _root;
            while(cur && cur->_left) {
                cur = cur->_left;
            }
            return iterator(cur, _root);
        }
        iterator end() { return iterator(nullptr, _root); }

        const_iterator begin() const
        {
            node* cur = _root;
            while(cur && cur->_left) {
                cur = cur->_left;
            }
            return const_iterator(cur,_root);
        }
        const_iterator end() const { return const_iterator(nullptr, _root); }

        const_iterator cbegin() const
        {
            node* cur = _root;
            while(cur && cur->_left) {
                cur = cur->_left;
            }
            return const_iterator(cur, _root);
        }
        const_iterator cend() const { return const_iterator(nullptr, _root); }
        /*TODO:修改insert函数，以及迭代器的operator--*/

    private:
        using node = rbtreeNode<Value>;
        void destroy(node* root)
        {
            if(root == nullptr) return;
            destroy(root->_left);
            destroy(root->_right);
            delete root;
            root = nullptr;
        }

        node* copy(const node* other)
        {
            if(other == nullptr) return nullptr;
            node* newnode = new node(other->_data);
            newnode->_col = other->_col;

            newnode->_left = copy(other->_left);
            newnode->_right = copy(other->_right);

            if(newnode->_left) newnode->_left->_parent = newnode;
            if(newnode->_right) newnode->_right->_parent = newnode;

            return newnode;
        }

    public:
        rbtree():_root(nullptr) {}
        ~rbtree() { destroy(_root); }
        rbtree(const rbtree& other) { _root = copy(other._root); }
        rbtree& operator=(rbtree other)
        {
            std::swap(_root, other._root);
            return *this;
        }

        // insert返回的是一个pair，iterator用来指向新插入元素的迭代器，或者已经存储的元素的迭代器
        // bool用来表示该元素是新插入的，还是原本已经有的，新插入的返回true，已经有的，返回false
        std::pair<iterator, bool> insert(const Value& data)
        {
            // 如果是空树，直接将新节点赋值给root
            if(_root == nullptr) {
                _root = new node(data);
                _root->_col = color::BLACK; // 根节点一定是黑色的
                return std::make_pair(iterator(_root, _root), true);
            }

            // 1.根据二叉树的插入逻辑，进行左右比较，插入到合适位置
            node* parent = nullptr;
            node* cur = _root;
            while(cur) {
                if(_kov(data) < _kov(cur->_data)) {
                    parent = cur;
                    cur = cur->_left;
                }else if(_kov(data) > _kov(cur->_data)) {
                    parent = cur;
                    cur = cur->_right;
                }else {
                    return std::make_pair(iterator(cur, _root), false); // 不允许冗余
                }
            }

            cur = new node(data);
            node* ret = cur; // 用来返回时构造迭代器，因为在后序调整中可能导致cur变化
            cur->_col = color::RED; // 新插入的节点除了根节点外必须是红色的
            if(_kov(data) < _kov(parent->_data)) {
                parent->_left = cur;
            }else {
                parent->_right = cur;
            }
            cur->_parent = parent;

            // 2.如果出现连续的红色，根据分类进行更新
            // 当g是根时，更新完成后不用再更新了，此时cur = root，parent = nullptr，此时直接退出，并且将根设置为黑色
            while(parent && parent->_col == color::RED) {
                node* grand = parent->_parent;

                if(grand->_left == parent) {
                    //       g
                    //    p     u
                    node* uncle = grand->_right;

                    // 当uncle存在且为红色，变色
                    // g变红，p、u变黑
                    if(uncle && uncle->_col == color::RED) {
                        parent->_col = color::BLACK;
                        uncle->_col = color::BLACK;
                        grand->_col = color::RED;

                        // 更新cur、parent，继续向上更新
                        cur = grand;
                        parent = cur->_parent;
                    }else {
                        // 当uncle不存在，或者存在为黑色
                        if(cur == parent->_left) {
                            //       g
                            //    p     u
                            // c
                            // 当一边高时，进行单旋+变色 :p变黑，g变红
                            rotateR(grand);
                            parent->_col = color::BLACK;
                            grand->_col = color::RED;
                        }else {
                            //       g
                            //    p     u
                            //      c
                            // 当出现右左不均时，进行双旋+变色：c变黑，g变红
                            rotateL(parent);
                            rotateR(grand);
                            cur->_col = color::BLACK;
                            grand->_col = color::RED;
                        }
                        break; // 旋转+变色之后不需要继续更新了
                    }
                }else {
                    //       g
                    //    u     p
                    node* uncle = grand->_left;

                    // uncle存在且为红，变色：u、p变黑，g变红
                    if(uncle && uncle->_col == color::RED) {
                        parent->_col = color::BLACK;
                        uncle->_col = color::BLACK;
                        grand->_col = color::RED;

                        // 更新cur、parent，继续向上更新
                        cur = grand;
                        parent = cur->_parent;
                    }else {
                        // uncle不存在/uncle存在为黑色，根据插入位置的不同选择单旋或者双旋，并且进行变色
                        if(cur == parent->_right) {
                            //       g
                            //    u     p
                            //             c
                            rotateL(grand);
                            // 单旋：p变黑，g变红
                            parent->_col = color::BLACK;
                            grand->_col = color::RED;
                        }else {
                            //       g
                            //    u     p
                            //        c
                            rotateR(parent);
                            rotateL(grand);
                            // 双旋：c变黑，g变红
                            cur->_col = color::BLACK;
                            grand->_col = color::RED;
                        }
                        break;
                    }
                }
            }
            _root->_col = color::BLACK; // 根必须是黑色的
            return std::make_pair(iterator(ret, _root), true);
        }

        node* find(const Value& data)
        {
            node* cur = _root;
            while(cur) {
                if(_kov(data) < _kov(cur->_data)) {
                    cur = cur->_left;
                }else if(_kov(data) > _kov(cur->_data)) {
                    cur = cur->_right;
                }else {
                    return cur;
                }
            }
            return nullptr;
        }
    private:
        // 右单旋
        void rotateR(node* prev)
        {
            //       pprev
            //       prev
            //   cur
            //     curR
            node* cur = prev->_left;
            node* curR = cur->_right;
            node* pprev = prev->_parent;

            //       pprev
            //       cur
            //            prev
            //          curR

            // 1.cur的右子树，变为prev的左子树
            if(curR) curR->_parent = prev; 
            prev->_left = curR;

            // 2. prev变为cur的右子树
            cur->_right = prev;
            prev->_parent = cur;

            // 3.更新cur的指向
            if(prev == _root) {
                _root = cur;
                cur->_parent = nullptr;
            }else {// 如果prev原先不是根，则判断cur应该连接到pprev的左子树还是右子树
                if(prev == pprev->_left) {
                    pprev->_left = cur;
                }else{
                    pprev->_right = cur;
                }
                cur->_parent = pprev;
            }
        }

        // 左单旋
        void rotateL(node* prev)
        {
            //   pprev
            //   prev
            //       cur
            //      curL
            node* cur = prev->_right;
            node* curL = cur->_left;
            node* pprev = prev->_parent;

            //      pprev
            //      cur
            //  prev
            //    curL

            // 1.cur的左子树，成为prev的右子树
            if(curL) curL->_parent = prev;
            prev->_right = curL;

            // 2.prev变为cur的左子树
            cur->_left = prev;
            prev->_parent = cur;

            // 3.更新cur的指向
            if(prev == _root) {
                _root = cur;
                cur->_parent = nullptr;
            }else {
                if(prev == pprev->_left) {
                    pprev->_left = cur;
                }else {
                    pprev->_right = cur;
                }
                cur->_parent = pprev;
            }
        }
    private:
        node* _root;
        KeyOfValue _kov;
    };
}